package com.dmall.graph.tree;

import com.dmall.graph.WeightedEdge;
import com.dmall.graph.WeightedGraph;
import com.dmall.graph.dfs.app.ConnectedComponent;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

/**
 * Prim算法寻找最小生成树
 * 优化后的Prim算法，使用了优先队列
 * @author xiao1.wang
 * @date 2022-08-05 09:56:00
 */
public class Prim2 {
    private WeightedGraph G;
    private List<WeightedEdge> mst;

    public Prim2(WeightedGraph G) {
        this.G = G;
        this.mst = new ArrayList<>(G.E());

        ConnectedComponent cc = new ConnectedComponent(G);
        if (cc.count() > 1) {
            return;
        }
        // Prim
        boolean[] visited = new boolean[G.V()];
        visited[0] = true;
        PriorityQueue<WeightedEdge> pq = new PriorityQueue<>();
        for (int w : G.adj(0)) {
            pq.add(new WeightedEdge(0, w, G.getWeight(0, w)));
        }
        while (!pq.isEmpty()) {
            WeightedEdge minEdge = pq.poll();
            if (visited[minEdge.getV()] && visited[minEdge.getW()]) {
                // 两个顶点都是已经访问过了
                continue;
            }
            mst.add(minEdge);
            int v = visited[minEdge.getV()] ? minEdge.getW() : minEdge.getV();
            visited[v] = true;
            for (int w : G.adj(v)) {
                if (!visited[w]) {
                    pq.offer(new WeightedEdge(v, w, G.getWeight(v, w)));
                }
            }
        }
    }

    public List<WeightedEdge> result() {
        return mst;
    }

    public static void main(String[] args) {
        String filename = Prim2.class.getClassLoader().getResource("wg.txt").getPath();
        WeightedGraph g = new WeightedGraph(filename);
        Prim2 prim = new Prim2(g);
        System.out.println(prim.result());
    }
}
